"""
通知服务
处理Webhook和WebSocket通知
"""

import asyncio
import json
import time
import hashlib
import hmac
from typing import Dict, Any, List, Optional, Set
import aiohttp
from fastapi import WebSocket

from ..storage import data_manager
from ..logger import logger
from ..config import config
from ..utils import retry_on_exception


class NotificationService:
    """通知服务"""
    
    def __init__(self):
        self.notification_config = config.get('notification', {})
        self.webhook_config = self.notification_config.get('webhook', {})
        self.websocket_config = self.notification_config.get('websocket', {})
        
        # WebSocket连接管理
        self.websocket_connections: Set[WebSocket] = set()
        self.client_connections: Dict[str, Set[WebSocket]] = {}
        
    async def start(self):
        """启动通知服务"""
        logger.info("Notification service started")
        
    async def stop(self):
        """停止通知服务"""
        # 关闭所有WebSocket连接
        for websocket in self.websocket_connections.copy():
            try:
                await websocket.close()
            except Exception:
                pass
        
        self.websocket_connections.clear()
        self.client_connections.clear()
        
        logger.info("Notification service stopped")
        
    async def send_webhook_notification(
        self, 
        url: str, 
        payload: Dict[str, Any],
        secret: Optional[str] = None
    ) -> bool:
        """发送Webhook通知"""
        try:
            if not self.webhook_config.get('enabled', True):
                logger.debug("Webhook notifications are disabled")
                return False
                
            headers = {
                "Content-Type": "application/json",
                "User-Agent": "ModelSubscriptionSystem/1.0"
            }
            
            # 添加签名
            if secret:
                signature = self._generate_webhook_signature(payload, secret)
                signature_header = config.get('security.webhook_secret_header', 'X-Webhook-Signature')
                headers[signature_header] = signature
            
            timeout = aiohttp.ClientTimeout(
                total=self.webhook_config.get('timeout', 30)
            )
            
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(url, json=payload, headers=headers) as response:
                    if response.status == 200:
                        logger.debug(f"Webhook notification sent successfully to {url}")
                        return True
                    else:
                        logger.warning(f"Webhook notification failed: {response.status} - {url}")
                        return False
                        
        except Exception as e:
            logger.error(f"Error sending webhook notification to {url}: {e}")
            return False
            
    def _generate_webhook_signature(self, payload: Dict[str, Any], secret: str) -> str:
        """生成Webhook签名"""
        payload_bytes = json.dumps(payload, sort_keys=True).encode('utf-8')
        signature = hmac.new(
            secret.encode('utf-8'),
            payload_bytes,
            hashlib.sha256
        ).hexdigest()
        return f"sha256={signature}"
        
    async def send_webhook_with_retry(
        self, 
        url: str, 
        payload: Dict[str, Any],
        secret: Optional[str] = None
    ) -> bool:
        """带重试的Webhook发送"""
        max_retries = self.webhook_config.get('retry_attempts', 3)
        
        for attempt in range(max_retries):
            success = await self.send_webhook_notification(url, payload, secret)
            if success:
                return True
                
            if attempt < max_retries - 1:
                # 指数退避
                delay = 2 ** attempt
                await asyncio.sleep(delay)
                
        return False
        
    async def broadcast_event(self, event_data: Dict[str, Any]):
        """广播事件到所有相关订阅者"""
        try:
            # 获取相关订阅
            subscriptions = await data_manager.get_subscriptions()
            
            # 过滤相关订阅
            relevant_subscriptions = []
            event_type = event_data.get('event_type', '')
            model_name = event_data.get('model_name', '')
            server_id = event_data.get('server_id', '')
            
            for subscription in subscriptions:
                if subscription.status != "active":
                    continue
                    
                # 检查模型匹配
                if model_name and subscription.model_name != model_name:
                    continue
                    
                # 检查服务器匹配
                if server_id and subscription.server_id != server_id:
                    continue
                    
                relevant_subscriptions.append(subscription)
            
            if not relevant_subscriptions:
                logger.debug(f"No relevant subscriptions for event: {event_type}")
                return
                
            logger.info(f"Broadcasting event {event_type} to {len(relevant_subscriptions)} subscriptions")
            
            # 发送通知
            tasks = []
            for subscription in relevant_subscriptions:
                # Webhook通知
                if "webhook" in subscription.notification_channels and subscription.callback_url:
                    task = asyncio.create_task(
                        self._send_subscription_webhook(subscription, event_data)
                    )
                    tasks.append(task)
                
                # WebSocket通知
                if "websocket" in subscription.notification_channels:
                    task = asyncio.create_task(
                        self._send_subscription_websocket(subscription, event_data)
                    )
                    tasks.append(task)
            
            # 等待所有通知完成
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)
                
        except Exception as e:
            logger.error(f"Error broadcasting event: {e}")
            
    async def _send_subscription_webhook(
        self, 
        subscription, 
        event_data: Dict[str, Any]
    ):
        """发送订阅Webhook通知"""
        try:
            payload = {
                "event": event_data.get('event_type', ''),
                "subscription_id": subscription.id,
                "client_id": subscription.client_id,
                "model_name": subscription.model_name,
                "server_id": subscription.server_id,
                "data": event_data,
                "timestamp": time.time()
            }
            
            success = await self.send_webhook_with_retry(subscription.callback_url, payload)
            
            if success:
                logger.debug(f"Webhook sent to subscription {subscription.id}")
            else:
                logger.warning(f"Failed to send webhook to subscription {subscription.id}")
                
        except Exception as e:
            logger.error(f"Error sending subscription webhook: {e}")
            
    async def _send_subscription_websocket(
        self, 
        subscription, 
        event_data: Dict[str, Any]
    ):
        """发送订阅WebSocket通知"""
        try:
            client_id = subscription.client_id
            
            if client_id not in self.client_connections:
                logger.debug(f"No WebSocket connections for client {client_id}")
                return
                
            message = {
                "type": "subscription_event",
                "event": event_data.get('event_type', ''),
                "subscription_id": subscription.id,
                "data": event_data,
                "timestamp": time.time()
            }
            
            # 发送到客户端的所有连接
            client_websockets = self.client_connections[client_id].copy()
            for websocket in client_websockets:
                try:
                    await websocket.send_text(json.dumps(message))
                except Exception as e:
                    logger.debug(f"Failed to send WebSocket message: {e}")
                    # 移除失效的连接
                    await self.remove_websocket_connection(websocket, client_id)
                    
        except Exception as e:
            logger.error(f"Error sending subscription WebSocket: {e}")
            
    async def add_websocket_connection(self, websocket: WebSocket, client_id: Optional[str] = None):
        """添加WebSocket连接"""
        try:
            self.websocket_connections.add(websocket)
            
            if client_id:
                if client_id not in self.client_connections:
                    self.client_connections[client_id] = set()
                self.client_connections[client_id].add(websocket)
                
            logger.debug(f"WebSocket connection added for client {client_id}")
            
        except Exception as e:
            logger.error(f"Error adding WebSocket connection: {e}")
            
    async def remove_websocket_connection(self, websocket: WebSocket, client_id: Optional[str] = None):
        """移除WebSocket连接"""
        try:
            self.websocket_connections.discard(websocket)
            
            if client_id and client_id in self.client_connections:
                self.client_connections[client_id].discard(websocket)
                
                # 如果客户端没有连接了，清理记录
                if not self.client_connections[client_id]:
                    del self.client_connections[client_id]
                    
            logger.debug(f"WebSocket connection removed for client {client_id}")
            
        except Exception as e:
            logger.error(f"Error removing WebSocket connection: {e}")
            
    async def send_websocket_message(
        self, 
        message: Dict[str, Any], 
        client_id: Optional[str] = None
    ):
        """发送WebSocket消息"""
        try:
            if not self.websocket_config.get('enabled', True):
                logger.debug("WebSocket notifications are disabled")
                return
                
            message_text = json.dumps(message)
            
            if client_id and client_id in self.client_connections:
                # 发送给特定客户端
                client_websockets = self.client_connections[client_id].copy()
                for websocket in client_websockets:
                    try:
                        await websocket.send_text(message_text)
                    except Exception as e:
                        logger.debug(f"Failed to send WebSocket message to client {client_id}: {e}")
                        await self.remove_websocket_connection(websocket, client_id)
            else:
                # 广播给所有连接
                websockets = self.websocket_connections.copy()
                for websocket in websockets:
                    try:
                        await websocket.send_text(message_text)
                    except Exception as e:
                        logger.debug(f"Failed to broadcast WebSocket message: {e}")
                        await self.remove_websocket_connection(websocket)
                        
        except Exception as e:
            logger.error(f"Error sending WebSocket message: {e}")
            
    async def get_connection_stats(self) -> Dict[str, Any]:
        """获取连接统计信息"""
        return {
            "total_connections": len(self.websocket_connections),
            "client_connections": {
                client_id: len(connections) 
                for client_id, connections in self.client_connections.items()
            },
            "webhook_enabled": self.webhook_config.get('enabled', True),
            "websocket_enabled": self.websocket_config.get('enabled', True)
        }


# 全局通知服务实例
notification_service = NotificationService()
